#!/usr/bin/env python3
"""
简化版客户购买预测脚本
"""
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.model_selection import KFold
from sklearn.metrics import roc_auc_score

print("=== 客户购买预测模型 ===")

# 读取数据
print("读取数据...")
train = pd.read_csv('train_set.csv')
test = pd.read_csv('test_set.csv')

print(f"训练集: {train.shape}, 测试集: {test.shape}")

# 特征工程
print("特征工程...")

# 数值特征
numeric_features = ['age', 'balance', 'day', 'duration', 'campaign', 'pdays', 'previous']

# 分类特征
categorical_features = ['job', 'marital', 'education', 'default', 'housing', 'loan', 'contact', 'month', 'poutcome']

# 处理分类特征 - one-hot编码
encoder = OneHotEncoder(handle_unknown='ignore', sparse_output=False)
categorical_data = pd.concat([train[categorical_features], test[categorical_features]])
encoder.fit(categorical_data)

train_encoded = pd.DataFrame(encoder.transform(train[categorical_features]))
test_encoded = pd.DataFrame(encoder.transform(test[categorical_features]))

train_encoded.columns = [f'cat_{i}' for i in range(train_encoded.shape[1])]
test_encoded.columns = [f'cat_{i}' for i in range(test_encoded.shape[1])]

# 合并特征
X = pd.concat([train[numeric_features], train_encoded], axis=1)
X_test = pd.concat([test[numeric_features], test_encoded], axis=1)
y = train['y']

print(f"特征形状: 训练{X.shape}, 测试{X_test.shape}")

# 标准化数值特征
scaler = StandardScaler()
X[numeric_features] = scaler.fit_transform(X[numeric_features])
X_test[numeric_features] = scaler.transform(X_test[numeric_features])

# 5折交叉验证训练
print("开始模型训练...")
n_folds = 5
kf = KFold(n_splits=n_folds, shuffle=True, random_state=42)

oof_preds = np.zeros(len(X))
test_preds = np.zeros(len(X_test))

for fold, (train_idx, val_idx) in enumerate(kf.split(X)):
    print(f"  训练第 {fold+1}/{n_folds} 折")
    
    X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
    y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]
    
    model = RandomForestRegressor(
        n_estimators=50,
        max_depth=8,
        random_state=42,
        n_jobs=-1
    )
    
    model.fit(X_train, y_train)
    oof_preds[val_idx] = model.predict(X_val)
    test_preds += model.predict(X_test) / n_folds

# 计算AUC分数
auc_score = roc_auc_score(y, oof_preds)
print(f"✅ 模型训练完成!")
print(f"📊 AUC 分数: {auc_score:.6f}")

# 生成预测文件
submission = pd.DataFrame({
    'ID': test['ID'],
    'pred': test_preds
})

output_file = f'submission_auc_{auc_score:.6f}.csv'
submission.to_csv(output_file, index=False)
print(f"📁 预测文件已保存: {output_file}")
print(f"📋 预测样本数量: {len(submission)}")

# 显示预测统计
print(f"📈 预测统计:")
print(f"   最小值: {test_preds.min():.4f}")
print(f"   最大值: {test_preds.max():.4f}")
print(f"   平均值: {test_preds.mean():.4f}")
print(f"   标准差: {test_preds.std():.4f}")

print("🎉 任务完成!")